import argparse
import torch
import os

from sentiment import *

if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('model' ,choices=['CNN','RNN', 'MLP'], help='which model to use')
    parser.add_argument('mode', choices=['train', 'eval', 'test'], help='train or eval')
    parser.add_argument('--gpu', action='store_true', help='use Nvidia gpu')
    parser.add_argument('--batch_size', action='store', type=int, default=16, help='batch size when train or eval [Default : 16]')
    parser.add_argument('--epoch', action='store', type=int, default=20, help='training epoches [Default : 20]')
    parser.add_argument('-s','--early_stop', action='store_true', help='use early stop when validation falls')
    parser.add_argument('-n', '--no_vector', action='store_true', help='DO NOT use pre trained word vector')
    parser.add_argument('-l', '--load_params', action='store', type=str, help='pre_trained model paramaters File Name')
    parser.add_argument('--path', action='store', type=str, help='path to the data set [Default : ./data]')
    parser.add_argument('--no_plot', action='store_true', help='Do NOT Plot the training figure')
    parser.add_argument('-w', '--weighted_sample', action='store_true', help='Use weighted sample')
    args = parser.parse_args()

    data_dir = os.path.abspath('./data') if not args.path else os.path.abspath(args.path)
    print('Build vocab...\r')
    train_dataset, val_dataset, test_dataset, TEXT = create_dataset(data_dir, validation=True, fixed_length=720,
                                                pretrained_vec=not args.no_vector)
    train_iter, val_iter, test_iter, weight = create_iterator(train_dataset, val_dataset, test_dataset, 
                                batch_size=args.batch_size, return_weight=True, weighted_sample=args.weighted_sample)
    print('...Build done!')
    # get device
    device = None
    if args.gpu:
        print("Use GPU")
        device = torch.device('cuda')
        weight = weight.to(device)
    # init net model
    model = None
    if args.model == 'CNN':
        config = CNNConfig(args.batch_size, pre_vector = not args.no_vector)
        model = textCNN(config, TEXT).to(device)
    elif args.model == 'RNN':
        config = RNNConfig(args.batch_size, pre_vector = not args.no_vector)
        model = textRNN(config, TEXT).to(device)
    elif args.model == 'MLP':
        config = MLPConfig(args.batch_size, pre_vector= not args.no_vector)
        model = MLP(config, TEXT).to(device)

    if args.load_params:
        if args.gpu:
            model.load_state_dict(torch.load(args.load_params))
        else:
            model.load_state_dict(torch.load(args.load_params, map_location=torch.device('cpu')))
    else:
        model.weight_init()

    # eval or train
    if args.mode == 'eval':
        evaluate(test_iter, model, data_name='test', device=device)
        exit(0)
    elif args.mode == 'test':
        exit(0)
    
    # optimizer = torch.optim.Adam(model.parameters(), lr=5e-3, weight_decay=1e-3)
    optimizer = torch.optim.Adam(model.parameters(), lr=5e-3, weight_decay=1e-3)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[10, 25, 50], gamma=0.2)
    val_acc, val_loss = train(model, train_iter, optimizer, val_it=val_iter, early_stop=args.early_stop,
                                weight=weight, device=device, epoches=args.epoch, scheduler=scheduler)
    if not args.no_plot:
        plot_val_acc_loss(val_acc, val_loss, model.__class__.__name__)
    
    evaluate(test_iter, model, data_name='test', device=device)